from datetime import datetime
from enum import Enum
from typing import Optional

from sqlalchemy import Text
from sqlalchemy.ext.asyncio import AsyncSession
from sqlmodel import Field, SQLModel, select


class TaskStatus(str, Enum):
    """
    任务状态
    """
    PENDING = "PENDING",
    SUCCESS = "SUCCESS",
    FAILURE = "FAILURE",
    EXECUTING = "EXECUTING",


class Task(SQLModel, table=True):
    """
     任务模型
    """
    __table_args__ = {'extend_existing': True}  # 允许重新定义表
    id: Optional[int] = Field(default=None, primary_key=True)
    input_path: str = Field(max_length=255, nullable=False)
    status: TaskStatus = Field(max_length=50, nullable=False)
    title: str = Field(max_length=500, nullable=True)
    content: str = Field(sa_type=Text, nullable=True)
    error_msg: str = Field(sa_type=Text, nullable=True)
    created_at: Optional[datetime] = Field(default_factory=datetime.now, nullable=False)
    start_execution_time: datetime = Field(nullable=True)
    end_execution_time: datetime = Field(nullable=True)


async def find_task_by_id(session: AsyncSession, task_id: int) -> Task | None:
    """
    根据id查询任务
    """
    statement = select(Task).where(Task.id == task_id)
    result = await session.execute(statement)
    return result.scalars().first()


async def save_task(session: AsyncSession, task: Task) -> Task:
    """
    保存任务
    """
    session.add(task)
    await session.commit()
    await session.refresh(task)
    return task
